import database
from queue import Queue

# input_pkg = 'Scrapy@1.8.2'
pkg_queue = Queue()
path = []

class RTree:
    def __init__(self, id):
        self.id = id
        self.children = []
        self.vuls = []

    def get_id(self):
        return self.id

    def add_child(self, node):
        self.children.append(node)

    def add_vuls(self, relations):
        for r in relations:
            vul = r.start_node
            if len(vul['CVE_ID']) > 0:
                self.vuls.append(vul['CVE_ID'])
            else:
                self.vuls.append(vul['ID'])

    # dfs遍历树并按先序打印路径
    def print_path(self):
        if len(self.vuls) == 0:
            path.append(self.id)
        else:
            str = self.id + '('
            for index in range(len(self.vuls)):
                if index < len(self.vuls) - 1:
                    str += self.vuls[index] + ','
                else:
                    str += self.vuls[index] + ')'
            path.append(str)
        if len(self.children) == 0:
            for index in range(len(path)):
                if index < len(path) - 1:
                    print(path[index], end = "->")
                else:
                    print(path[index])
            print()
            path.pop()
            return

        for child in self.children:
            child.print_path()

        path.pop()
        return

# bfs方式搜图建树
def bfs(id):
    root = RTree(id)
    pkg_queue.put(root)
    db = database.DataBase()
    while not pkg_queue.empty():
        cur = pkg_queue.get()
        node1 = db.match_pkg_by_id(cur.get_id())
        if node1 is None:
            print("%s is not in database" % cur.get_id())
            break
        vuls = db.query_relations_by_node2(node1, 'affect_pkg')
        if len(vuls) != 0:
            cur.add_vuls(vuls)
        rela_list = db.query_relations_by_node1(node1, 'depends_on_pkg')
        for relation in rela_list:
            node2 = relation.end_node
            nxt = RTree(node2['ID'])
            cur.add_child(nxt)
            pkg_queue.put(nxt)
    return root

if __name__ == '__main__':
    print("please input the pkg-id which you are searching for:")
    input_pkg = input()
    tree = bfs(input_pkg)
    tree.print_path()


